# coding:utf-8
# @Time    : 2021/11/10 15:20
# @Author  : Martin
from selenium.webdriver.support.wait import WebDriverWait, POLL_FREQUENCY
import time
from selenium.common.exceptions import InvalidSelectorException
from selenium.common.exceptions import TimeoutException

"""
WebDriverWait扩展类
重写util方法，使其返回TimeoutException时打印具体的异常信息
"""


class WebDriverWaitExtend(WebDriverWait):
    def __init__(self, driver, timeout, poll_frequency=POLL_FREQUENCY, ignored_exceptions=None):
        super(WebDriverWaitExtend, self).__init__(driver, timeout, poll_frequency, ignored_exceptions)

    def until(self, method, message=None):
        """Calls the method provided with the driver as an argument until the \
        return value does not evaluate to ``False``.

        :param method: callable(WebDriver)
        :param message: optional message for :exc:`TimeoutException`
        :returns: the result of the last call to `method`
        :raises: :exc:`selenium.common.exceptions.TimeoutException` if timeout occurs
        """
        screen = None
        stacktrace = None

        end_time = time.time() + self._timeout
        while True:
            try:
                value = method(self._driver)
                if value:
                    return value
            except InvalidSelectorException as e:
                raise e
            except self._ignored_exceptions as exc:
                # screen = getattr(exc, 'screen', None)
                # stacktrace = getattr(exc, 'stacktrace', None)
                # 新增错误异常信息
                message += getattr(exc, 'msg', None)
            time.sleep(self._poll)
            if time.time() > end_time:
                break
        raise TimeoutException(message, screen, stacktrace)
